# -*- encoding: utf-8 -*-
"""
@File    :   algorithm.py
@Contact :   haoyu_deng@std.uestc.edu.cn


@Modify Time      @Author          @Version    @Desciption
------------      -------------    --------    -----------
2022/11/20 18:57   risingentropy      1.0         None
"""

from utils.tensorbackend import get_backend, BackendBase


class AlgorithmBase:
    """
    Base class for all algorithms
    Attributes:
        weight: the filtering weight, if not fitted, remain None. The shape of weight is automatically created and will be randomly initialized in the first fitting.
        That is to say you don't have to deal with it if there's no need for convenience
    """
    weight = None

    def __init__(self):
        pass

    def iterate(self, target, input, **extArgs):
        raise NotImplemented("AlgorithmBase is an abstract class and its methods cannot be invoked")


class ObserveOverTargetAlgorithmBase(AlgorithmBase):
    """
        Filtering in the situation that only target is observed, like LMS RLS algorithms.
    """

    def __init__(self):
        """
        """
        super().__init__()
        self.n = None

    def iterate(self, target, input, **extArgs):
        """

        :param target:
        :param input:
        :param extArgs:
        :return:
        """
        backend = get_backend(input)

        if self.n is None:
            self.n = backend.shape(input)[0]
        elif self.n != backend.shape(input)[0]:
            raise ValueError(
                f"the input data length{backend.shape(input)[0]} doesn't match the former weight shape{self.n}")

        if self.weight is None:
            self.weight = backend.normal(loc=0, scale=1, shape=(self.n, 1))

        self._iterate(target, input, backend, extArgs)
        return self.weight

    def _iterate(self, target, input, backend, extArgs):
        """

        :param target:
        :param input:
        :param extArgs:
        :return:
        """
        raise NotImplementedError(
            "ObserveOverTargetAlgorithmBase is an abstract class and its methods cannot be invoked")


class LMS(ObserveOverTargetAlgorithmBase):
    """

    """

    def __init__(self, step_size: float):
        """

        :param step_size:
        """
        super().__init__()
        self.step_size = step_size

    def _iterate(self, target, input, backend, extArgs):
        """

        :param target:
        :param input:
        :param extArgs:
        :return:
        """

        for k in range(0, input.shape[-1]):  # 默认最后一维度为数据长度
            err = target[:, k] - backend.matmul(backend.transpose(self.weight), backend.unsqueeze(input[:, k], dim=-1))
            self.weight = self.weight + self.step_size * err * backend.reshape(input[:, k], backend.shape(self.weight))


class NLMS(ObserveOverTargetAlgorithmBase):
    def __init__(self, step_size: float):
        """

        :param step_size:
        """
        super().__init__()
        self.step_size = step_size

    def _iterate(self, target, input, backend, extArgs):
        """

        :param target:
        :param input:
        :param extArgs:
        :return:
        """
        for k in range(0, input.shape[-1]):  # 默认最后一维度为数据长度
            err = target[:, k] - backend.matmul(backend.transpose(self.weight), backend.unsqueeze(input[:, k], dim=-1))
            self.weight = self.weight + backend.reshape(self.step_size * err * input[:, k] / backend.matmul(
                backend.transpose(backend.unsqueeze(input[:, k], dim=-1)), backend.unsqueeze(input[:, k], dim=-1)),
                                                        backend.shape(self.weight))


class MCC(ObserveOverTargetAlgorithmBase):
    def __init__(self, step_size: float, sigma: float = 2):
        """

        :param step_size:
        :param sigma
        """
        super().__init__()
        self.step_size = step_size
        self.sigma = sigma

    def _iterate(self, target, input, backend, extArgs):
        for k in range(0, input.shape[-1]):  # 默认最后一维度为数据长度
            err = target[:, k] - backend.matmul(backend.transpose(self.weight), backend.unsqueeze(input[:, k], dim=-1))
            self.weight = self.weight + backend.reshape(
                self.step_size * backend.exp(-(err ** 2) / (2 * self.sigma ** 2)) * err * input[:, k],
                backend.shape(self.weight))

class MEE(ObserveOverTargetAlgorithmBase):
    def __init__(self, step_size: float, sigma: float = 2):
        """

        :param step_size:
        :param sigma
        """
        super().__init__()
        self.step_size = step_size
        self.sigma = sigma

    def _iterate(self, target, input, backend, extArgs):
        for k in range(0, input.shape[-1]):  # 默认最后一维度为数据长度
            err = target[:, k] - backend.matmul(backend.transpose(self.weight), backend.unsqueeze(input[:, k], dim=-1))
            self.weight = self.weight + backend.reshape(
                self.step_size * backend.exp(-(err ** 2) / (2 * self.sigma ** 2)) * err * input[:, k],
                backend.shape(self.weight))

class RLS(ObserveOverTargetAlgorithmBase):
    """

    """
    P = None

    def __init__(self, step_size: float):
        """

        :param step_size: treat as lambda
        """
        super().__init__()
        self.labda = 1
        self.step_size = step_size

    def iterate(self, target, input, **extArgs):
        """

        :param target:
        :param input:
        :param extArgs: 说明一下这里面可以有个P指定，不指定为I  lambda指定，不指定默认为1
        :return:
        """
        backend = get_backend(input)
        if self.n is None:
            self.n = backend.shape(input)[0]
        elif self.n != backend.shape(input)[0]:
            raise ValueError(
                f"the input data length{backend.shape(input)[0]} doesn't match the former weight shape{self.n}")

        if self.weight is None:
            self.weight = backend.zeros(shape=(self.n, 1))

        return super().iterate(target, input, **extArgs)

    def _iterate(self, target, input, backend, extArgs):
        """

        :param target:
        :param input:
        :param backend:
        :param extArgs:
        :return:
        """
        for k in range(0, input.shape[-1]):  # 默认最后一维度为数据长度
            err = target[:, k] - backend.matmul(self.weight.T, backend.unsqueeze(input[:, k], axis=1))
            K = backend.matmul(self.P, backend.expand_dims(input[:, k], axis=1)) / (
                    self.labda + backend.matmul(backend.unsqueeze(input[:, k], dim=1).T,
                                                backend.matmul(self.P, backend.unsqueeze(input[:, k], dim=1))))
            self.P = 1. / self.labda * (
                    self.P - backend.matmul(K, backend.matmul(backend.unsqueeze(input[:, k], dim=1).T, self.P)))
            self.weight = self.weight + backend.matmul(K, err)


class KalmanFiltering(AlgorithmBase):
    def __init__(self, A_k, C_k, Q_k, R_k):
        super().__init__()
        self.A_k = A_k
        self.C_k = C_k
        self.Q_k = Q_k
        self.R_k = R_k
        self.P = None

    def updateQ(self, Q):
        self.Q_k = Q

    def updateA(self, A):
        self.A_k = A

    def updateC(self, C):
        self.C_k = C

    def updateR(self, R):
        self.R_k = R

    def iterate(self, target, input, **extArgs):
        """

        :param target:
        :param input:
        :param extArgs:
        :return:
        """
        backend = get_backend(input)
        if self.P is None:
            self.P = backend.eye(backend.shape(input)[0], device=backend.device(input))

        return self._iterate(target, input, backend, extArgs)

    def _iterate(self, target, input, backend: BackendBase, extArgs):
        self.P = backend.matmul(backend.matmul(self.A_k, self.P), backend.transpose(self.A_k)) + self.Q_k
        H = backend.matmul(backend.matmul(self.P, backend.transpose(self.C_k)),
                           backend.inv(
                               backend.matmul(backend.matmul(self.C_k, self.P),
                                              backend.transpose(self.C_k)) + self.R_k))
        x_predic = backend.matmul(self.A_k, input) + backend.matmul(H, target - backend.matmul(
            backend.matmul(self.C_k, self.A_k), input))
        self.P = backend.matmul(
            backend.eye(backend.shape(input)[0], device=backend.device(input)) - backend.matmul(H, self.C_k), self.P)

        return x_predic
